package com.ruoyi.web.controller.chat;

import com.alibaba.fastjson.JSONObject;
import com.ruoyi.web.controller.chat.util.ChangeDTO;
import com.ruoyi.web.core.config.ChatModel;
import com.ruoyi.web.core.config.HfThreadLocal;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.MediaType;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;

/**
 * TODO
 *
 * @author chenchuan2
 * @date 2023-05-24 17:54:56
 */
@Slf4j
public class HfSseEmitter extends SseEmitter {
    @Override
    public void send(Object object, MediaType mediaType) throws IOException {
        log.info("stream回答问题[{}]",object);
        ChangeDTO model = HfThreadLocal.getModel();
        if (!object.toString().startsWith("{")) {
            super.send(object, mediaType);
        } else {
            if (model.getChange()) {
                JSONObject jsonObject = JSONObject.parseObject(object.toString());
                if (jsonObject.containsKey("model")) {
                    jsonObject.put("model", ChatModel.queryByModelName(model.getRealModelName()).getModelLongName());
                    super.send(jsonObject, mediaType);
                } else {
                    super.send(object, mediaType);
                }
            }else{
                super.send(object, mediaType);
            }
        }
    }
}
